"""
基于卫星测高数据获取的绝对动力地形数据（Absolute Dynamic Topography，ADT），
以及风场（Wind）数据，提取全球分辨率为1度的月平均或日平均海表流场数据。
"""

import glob
from collections import deque
from datetime import datetime
import math
import os
import sys
import time
import numpy as np
import pandas as pd
import netCDF4 as nc
from tqdm import tqdm  
import calendar

root_path = r"__file__"                                                                 # 添加pydggs路径
if root_path not in sys.path:
    sys.path.insert(0, root_path)

from PyDGGS.PyDGGS import *
from common import NAN, TIFF_NODATA_VALUE, level, grid_type, element_type               # 导入常量

time_resolution = 'daily'                                                               # 设置时间分辨率:   月：monthly，日：daily

input_dir = os.path.join(os.path.dirname(r'__file__'))                                  # 输入文件目录    
ocean_grid_path = os.path.join(input_dir, 'output', 'Ocean_Grid','Ocean_Grid.json')     # 海洋格网文件

# 设置daily的工作路径
daily_input_adt_path = os.path.join(input_dir, 'input', 'ADT_daily')                    # ADT文件
daily_input_wind_path = os.path.join(input_dir, 'input', 'Wind_daily')                  # Wind文件
daily_temp_output_path = os.path.join(input_dir, 'temp', 'daily_uv')                    # 临时工作目录，存储从原始nc文件中提取数据到以格网编码为索引的CSV数据文件。
daily_uv_output_path = os.path.join(input_dir, 'output', 'daily_uv')                    # 存储提取出的流场数据

# 设置monthly的工作路径
monthly_input_adt_path = os.path.join(input_dir, 'input', 'ADT_monthly')                # ADT文件
monthly_input_wind_path = os.path.join(input_dir, 'input', 'Wind_monthly')               # Wind文件
monthly_temp_output_path = os.path.join(input_dir, 'temp', 'monthly_uv')                # 临时工作目录，存储从原始nc文件中提取数据到以格网编码为索引的CSV数据文件。
monthly_uv_output_path = os.path.join(input_dir, 'output', 'monthly_uv')                # 存储提取出的流场数据


def extractNcData(var_file, var, year, month, day):
    """
    从NetCDF文件中提取数据

    :param var_file:    NetCDF路径或目录
    :param var:         属性
    :param year:        年
    :param month:       月
    :param day:         日（可选）
    :return:            属性值List, 每个值所在点的经纬度)
    """
    # 读取不同时间分辨率的nc文件
    if os.path.isdir(var_file):
        if day != '':
            date_str = f"{year}{month.zfill(2)}{day.zfill(2)}"
            file_pattern = f"*{date_str}*.nc"
        else:
            date_str = f"{year}{month.zfill(2)}"
            file_pattern = f"*{date_str}*.nc"

        matching_files = glob.glob(os.path.join(var_file, file_pattern))
        
        if not matching_files:
            if day != '':
                date_str = f"{year}{month.zfill(2)}"
                file_pattern = f"*{date_str}*.nc"
                matching_files = glob.glob(os.path.join(var_file, file_pattern))
            
            if not matching_files:
                raise FileNotFoundError(f"未找到匹配 {date_str} 的{var}文件")
        
        var_file = matching_files[0]
        print(f"使用文件: {os.path.basename(var_file)}")

    # 读取每天的nc文件
    if day != '':
        year = int(year)
        month = int(month)
        day = int(day)

        with open(var_file, mode='rb') as f:
            with nc.Dataset(var_file, mode='r', memory=f.read()) as ncfile:
                
                if "wind" in var_file.lower() and len(ncfile.variables['time']) > 1:

                    target_date_start = datetime(year, month, day, 0, 0, 0)
                    target_date_end = datetime(year, month, day, 23, 59, 59)
                    # 获取时间变量
                    time_var = ncfile.variables['time']
                    time_units = time_var.units
                    time_calendar = time_var.calendar

                    time_values = time_var[:]
                    time_dates = nc.num2date(time_values, units=time_units, calendar=time_calendar)

                    daily_indices = []
                    for idx, t in enumerate(time_dates):
                        if target_date_start <= t <= target_date_end:
                            daily_indices.append(idx)

                    if not daily_indices:
                        print(f"警告: {target_date_start} 无数据，使用最近时间点")
                        time_diffs = [abs(t - target_date_start) for t in time_dates]
                        closest_idx = np.argmin(time_diffs)
                        daily_indices = [closest_idx]

                    var_data = ncfile.variables[var]
                    if len(var_data.shape) == 3:  
                        daily_data = np.mean(var_data[daily_indices, :, :], axis=0)
                    else:
                        daily_data = var_data[:]

                    lat = ncfile.variables['latitude'][:]
                    lon = ncfile.variables['longitude'][:]
                    var_result = []

                    # 显示进度条
                    total_points = lat.size * lon.size
                    print(f"提取风场数据: {var}, 共 {total_points} 个点")

                    for lat_idx in tqdm(range(lat.size), desc=f"处理纬度点"):
                        for lon_idx in range(lon.size):
                            temp_lat = lat[lat_idx]
                            temp_lon = lon[lon_idx]
                            value = daily_data[lat_idx, lon_idx]

                            if temp_lon > 180:
                                temp_lon = temp_lon - 360.0

                            var_result.append((value, temp_lon, temp_lat))

                    lon_result = [pt[1] for pt in var_result]
                    lat_result = [pt[2] for pt in var_result]
                    return var_result, lon_result, lat_result

                variable = ncfile.variables[var]
                lat = ncfile.variables['latitude'][:]
                lon = ncfile.variables['longitude'][:]

                # 获取 time 变量
                time_var = ncfile.variables['time']
                if np.ma.is_masked(time_var[:]):
                    time_values = np.ma.compressed(time_var[:])
                else:
                    time_values = time_var[:]

                # 通过指定年月日获取属性索引
                target_date = datetime(year, month, day)
                time_dates = nc.num2date(time_values, units=time_var.units, calendar=time_var.calendar)
                closest_time_index = (abs(time_dates - target_date)).argmin()
                closest_time = time_dates[closest_time_index]

                # 计算索引
                time_diff = closest_time - target_date
                total_days_diff = time_diff.days
                time_index = closest_time_index  

                # 提取非掩码数据及其对应的 lat 和 lon
                if len(variable.shape) == 4:
                    var_data = variable[time_index, 0, :, :]
                elif len(variable.shape) == 3:
                    var_data = variable[time_index, :, :]
                else:
                    var_data = variable[:]

                var_data_array = np.ma.array(var_data)
                non_masked_indices = np.where(~var_data_array.mask)

                lat_indices, lon_indices = non_masked_indices[0], non_masked_indices[1]
                var_result = []
                lat_result = []
                lon_result = []

                # 添加进度条
                total_points = len(lat_indices)
                print(f"提取ADT数据: {var}, 共 {total_points} 个点")

                for idx in tqdm(range(len(lat_indices)), desc="处理点"):
                    lat_idx, lon_idx = lat_indices[idx], lon_indices[idx]
                    temp_lat = lat[lat_idx]
                    temp_lon = lon[lon_idx]
                    depth = var_data_array[lat_idx, lon_idx]

                    if temp_lon > 180:
                        temp_lon = temp_lon - 360.0

                    lat_result.append(temp_lat)
                    lon_result.append(temp_lon)
                    var_result.append((depth, temp_lon, temp_lat))

        return var_result, lon_result, lat_result

    else:
        year = int(year)
        month = int(month)

        with open(var_file, mode='rb') as f:
            with nc.Dataset(var_file, mode='r', memory=f.read()) as ncfile:
                variable = ncfile.variables[var]
                lat = ncfile.variables['latitude'][:]
                lon = ncfile.variables['longitude'][:]

                # 获取 time 变量
                time_var = ncfile.variables['time']
                if np.ma.is_masked(time_var[:]):
                    time_values = np.ma.compressed(time_var[:])  # 压缩掩盖的值
                else:
                    time_values = time_var[:]

                # 通过指定年月获取属性索引
                target_date = datetime(year, month, 1)
                time_dates = nc.num2date(time_values, units=time_var.units, calendar=time_var.calendar)
                closest_time_index = (abs(time_dates - target_date)).argmin()
                closest_time = time_dates[closest_time_index]

                year_diff = closest_time.year - year
                month_diff = closest_time.month - month
                total_month_diff = year_diff * 12 + month_diff
                time_index = variable.shape[0] - total_month_diff - 1

                # 提取非掩码数据及其对应的 lat 和 lon
                if len(variable.shape) == 4:
                    var_data = variable[time_index, 0, :, :]
                else:
                    var_data = variable[0, :, :]
                var_data_array = np.ma.array(var_data)
                non_masked_indices = np.where(~var_data_array.mask)

                lat_indices, lon_indices = non_masked_indices[0], non_masked_indices[1]
                var_result = []
                lat_result = []
                lon_result = []

                # 获取属性及其经纬度
                for lat_idx, lon_idx in zip(lat_indices, lon_indices):
                    temp_lat = lat[lat_idx]
                    temp_lon = lon[lon_idx]
                    depth = var_data_array[lat_idx, lon_idx]
                    if temp_lon > 180:
                        temp_lon = temp_lon - 360.0

                    lat_result.append(temp_lat)
                    lon_result.append(temp_lon)
                    var_result.append((depth, temp_lon, temp_lat))

        return var_result, lon_result, lat_result


def getGridCoordinates(level, grid_type, element_type):
    """
    根据指定等级，指定格网类型的全球离散格网坐标

    :param level:           离散格网等级
    :param grid_type:       离散格网类型
    :param element_type:    离散格网元素类型
    :return:                指定等级、格网类型对应的全球离散格网坐标
    """
    py_kernel = CreateKernelObj(grid_type)

    # 计算id
    max_row = rowMaxSize(level)
    max_col = fColMaxSize(level, grid_type)
    if element_type == DGGSElementType_Edge:
        max_col = eColMaxSize(level, grid_type)
    elif element_type == DGGSElementType_Vertex:
        max_col = vColMaxSize(level, grid_type)

    # 计算经纬度
    coordinates = []
    for base_partition in range(1,11):
        for row in range(max_row):
            for col in range(max_col):
                code = DGGSElementCode()
                code.basePartition = base_partition
                code.level = level
                code.row = row
                code.col = col
                code.elementType = element_type

                coord = GeoCoord()
                py_kernel.decode(code, coord)

                point = {}
                code_id = code2Id(code2Id, element_type)
                point['id'] = code_id
                point['lon'] = coord.lon
                point['lat'] = coord.lat
                coordinates.append(point)

    return coordinates


def createIndexCsv(var, year, month, day, input_file, index_path, 
                     level, grid_type, element_type, ocean_codeid_list):
    """
    存储从原始nc文件中提取数据到以格网编码为索引的CSV数据文件。

    :param var:                  属性名
    :param year:                 年份
    :param month:                月份
    :param day:                  日期
    :param input_file:           NetCDF路径
    :param index_path:           索引文件csv的输出路径
    :param level:                离散格网等级
    :param grid_type:            离散格网类型
    :param element_type:         离散格网元素类型
    :param ocean_codeid_list:    海洋格网的codeID列表
    """
    # 根据时间分辨率生成 layer_name 与 current_time
    if day != '':
        current_time = pd.to_datetime(f'{year}-{month}-{day}')
        out_csv = os.path.join(index_path, f"{year}{month}{day}_{var}.csv")
    else:
        current_time = pd.to_datetime(f'{year}-{month}-15')
        out_csv = os.path.join(index_path, f"{year}{month}_{var}.csv")

    # 提取原始数据
    var_data, _, _ = extractNcData(input_file, var, year, month, day)

    # 先映射到格网
    py_kernel = CreateKernelObj(grid_type)
    values = []
    for val, lon, lat in var_data:
        coord = GeoCoord()
        coord.lon = float(lon)
        coord.lat = float(lat)
        code = DGGSElementCode()
        code.level = level
        code.elementType = element_type
        py_kernel.encode(coord, code)
        cid = code2Id(code, grid_type)

        # 异常值剔除
        if var == 'adt' and abs(val) > 1.5:
            val = 0.00001
        elif var in ['eastward_wind', 'northward_wind'] and abs(val) > 20:
            val = 0.00001
        values.append((cid, val))

    # 插值4次
    for _ in range(4):
        d = {cid: v for cid, v in values}
        new_values = []

        for cid in ocean_codeid_list:
            v = d.get(cid, NAN)
            if v == NAN:
                # 重新生成 code，再求邻居
                code = DGGSElementCode()
                code.level = level
                code.elementType = element_type
                id2Code(cid, code, grid_type)

                nb_list = findInterRange(code, ocean_codeid_list, grid_type)
                total, cnt = 0.0, 0
                for nb in nb_list:
                    nb_v = d.get(nb, NAN)
                    if nb_v != NAN:
                        total += nb_v
                        cnt += 1
                if cnt:
                    v = total / cnt
            if v != NAN:
                new_values.append((cid, v))
        values = new_values

    DestroyKernelObj(py_kernel)

    # 输出 CSV 
    df = pd.DataFrame(values, columns=['id', var])
    df.insert(1, 'datetime', current_time)   # 保持列顺序
    df.to_csv(out_csv, index=False)


def createIndexPath(data_vars, year, month, day, 
                      input_path, temp_path, level, grid_type, element_type, ocean_code_id_list):
    """
    生成索引文件夹目录

    :param data_vars:           属性列
    :param year:                年
    :param month:               月
    :param day:                 日
    :param input_path:          输入文件夹或文件
    :param temp_path:           临时文件夹
    :param level:               格网等级
    :param grid_type:           格网类型
    :param element_type:        元素类型
    :param ocean_code_id_list:  海洋格网ID列表
    """
    os.makedirs(temp_path, exist_ok=True)
    if day != '':
        for var in data_vars:
            input_file = input_path
            createIndexCsv(
                var, year, month, day,
                input_file, temp_path,
                level, grid_type, element_type, ocean_code_id_list
            )

    else:
        is_file = os.path.isfile(input_path)
        file_pattern = f"dt_global_allsat_madt_h_y{year}_m{month}.nc"
    
        for var in data_vars:
            if is_file:
                input_file = input_path
            else:
                input_file = os.path.join(input_path, file_pattern)
            createIndexCsv(
                var, year, month, day,
                input_file, temp_path,
                level, grid_type, element_type, ocean_code_id_list
            )


def calOceanFlow(year, month, day, 
                   input_files, output_file_path, level, grid_type):
    """
    计算指定日期的流场

    :param year:               年份
    :param month:              月份
    :param day:                日期
    :param input_files:        索引文件路径元组 (东方向风场，北方向风场，ADT)
    :param output_file_path:   流场输出路径
    :param level:              离散格网等级
    :param grid_type:          离散格网类型
    """
    east_wind_file = input_files[0]
    north_wind_file = input_files[1]
    adt_file = input_files[2]

    print(f"\n计算流场: {year}-{month}-{day}")

    east_wind_data = pd.read_csv(east_wind_file)
    north_wind_data = pd.read_csv(north_wind_file)
    adt_data = pd.read_csv(adt_file)

    merged_data = pd.merge(east_wind_data, north_wind_data, on=['id', 'datetime'], suffixes=('_east', '_north'))
    merged_data = pd.merge(merged_data, adt_data, on=['id', 'datetime'])

    # 重命名列
    merged_data.rename(columns={
        'eastward_wind': 'east_wind',
        'northward_wind': 'north_wind',
        'adt': 'adt'
    }, inplace=True)

    wind_and_adt = merged_data[['id', 'datetime', 'east_wind', 'north_wind', 'adt']]
    cell_id_list = wind_and_adt['id'].tolist()
    id_to_adt = dict(zip(wind_and_adt['id'], wind_and_adt['adt']))

    # 对每个格元进行地转流计算
    py_kernel = CreateKernelObj(grid_type)

    band_count = 4
    element_type = DGGSElementType_Cell
    field_layer = DGGSFieldLayerDouble(
        band_count, DGGS_Float64, DGGSElementType_Cell, grid_type, level
    )
    row_size = rowMaxSize(level)
    col_size = fColMaxSize(level, grid_type)

    # 初始化格网
    for base_partition in range(1, 11):
        for row in range(row_size):
            for col in range(col_size):
                for band_id in range(band_count):
                    code = DGGSElementCode()
                    code.level = level
                    code.basePartition = base_partition
                    code.elementType = element_type
                    code.row = row
                    code.col = col
                    field_layer.setValue(code, TIFF_NODATA_VALUE, band_id)

    total_cells = len(wind_and_adt)
    pbar = tqdm(total=total_cells, desc="计算格元")

    for idx, row in wind_and_adt.iterrows():
        cell_id = int(row['id'])
        code = DGGSElementCode()
        code.level = level
        code.elementType = DGGSElementType_Cell
        id2Code(cell_id, code, grid_type)

        if not (0 < code.basePartition < 11 and 0 <= code.row <= row_size and 0 <= code.col <= col_size):
            pbar.update(1)
            continue

        coord = GeoCoord()
        py_kernel.decode(code, coord)
        lon = coord.lon
        lat = coord.lat

        adt = float(row['adt'])
        eastward_wind = float(row['east_wind'])
        northward_wind = float(row['north_wind'])

        # 计算结果变量
        geos_flow_x = NAN
        geos_flow_y = NAN
        ekman_flow_x = NAN
        ekman_flow_y = NAN
        total_flow_x = NAN
        total_flow_y = NAN
        total_flow_v = NAN
        total_flow_angle = NAN

        has_valid_neighbors = True
        query_result = QueryResult()
        py_kernel.query(DGGSTopoType_RelativeCells, code, query_result)
        cell_num = query_result.getSize()

        # 检查邻域格网是否有效
        for i in range(cell_num):
            neighbour_cell = query_result.getGridCodeByIndex(i)
            neighbour_cell_id = code2Id(neighbour_cell, grid_type)
            if neighbour_cell_id not in cell_id_list:
                has_valid_neighbors = False
                break

        if has_valid_neighbors:
            # 地转流的计算
            f = calCoriolis(lat)
            if f != 0:
                # 以当前格元中心为局部坐标中心（0,0）
                X = [0]
                Y = [0]
                Z = [adt]

                # 将该格元周边相邻格元的局部坐标以及adt值存入
                for i in range(cell_num):
                    coord = GeoCoord()
                    neighbour_cell = query_result.getGridCodeByIndex(i)
                    py_kernel.decode(neighbour_cell, coord)
                    coordinate_x, coordinate_y = latlon2Coordinate(lat, lon, coord.lat, coord.lon)
                    X.append(coordinate_x)
                    Y.append(coordinate_y)

                    # 获取当前邻接格元的adt
                    neighbour_cell_id = code2Id(neighbour_cell, grid_type)
                    neighbour_adt = id_to_adt.get(neighbour_cell_id, NAN)
                    Z.append(float(neighbour_adt))

                # 使用二次方程拟合，返回导数值
                differential = getDifferential2(X, Y, Z, 0, 0, "all")

                # 计算 f 平面近似计算地转流
                g = 9.7803 * (1 + 0.0053024 * (math.sin(math.radians(lat))) ** 2 - 0.000005 * math.sin(
                    math.radians(2 * lat)))
                geos_flow_x = -g / f * differential[1]  # x
                geos_flow_y = g / f * differential[0]  # y

                # 在低纬度(5°S-5°N)地区内，用β平面近似计算地转流
                if -5 < lat < 5:
                    C = 0.7  # weighting factor
                    theta = 0.2  # 纬向长度尺度
                    WB = C * math.exp(-((lat / theta) ** 2))
                    Wf = 1 - math.exp(-((lat / theta) ** 2))
                    geos_flow_x_b = -g / (2.3e-11) * differential[2]  # yy
                    geos_flow_y_b = g / (2.3e-11) * differential[3]  # xy
                    geos_flow_x = Wf * geos_flow_x + WB * geos_flow_x_b
                    geos_flow_y = Wf * geos_flow_y + WB * geos_flow_y_b

                # 进行Ekman流计算
                # 计算综合风速
                wind_speed = (northward_wind ** 2 + eastward_wind ** 2) ** 0.5
                swaying = calSwayingCoefficient(wind_speed)
                pa_air = 1.2  # 空气密度(1.2kg/m3)
                tx = eastward_wind * pa_air * swaying * wind_speed
                ty = northward_wind * pa_air * swaying * wind_speed

                # 计算Ekman流 vx vy
                if abs(lat) < 10:
                    ekman_flow_x, ekman_flow_y = calEquatorEkman(lat, tx, ty)
                elif abs(lat) <= 25:
                    weight10 = math.exp(-((abs(lat) - 10) / 6.6) ** 2)
                    weight25 = 1 - weight10

                    ekman_equator = calEquatorEkman(lat, tx, ty)
                    B = 0.3  # 中间参数

                    if lat >= 0:
                        O = -55.0 / 180.0 * math.pi  # 偏右
                    else:
                        O = 55.0 / 180.0 * math.pi  # 偏左

                    vxy25 = B * (complex(math.cos(O), math.sin(O))) * complex(tx, ty)
                    vx25 = vxy25.real
                    vy25 = vxy25.imag

                    ekman_flow_x = weight10 * ekman_equator[0] + weight25 * vx25
                    ekman_flow_y = weight10 * ekman_equator[1] + weight25 * vy25
                else:
                    B = 0.3  # 中间参数

                    if lat >= 0:
                        O = -55.0 / 180.0 * math.pi  # 偏右
                    else:
                        O = 55.0 / 180.0 * math.pi  # 偏左

                    vxy = B * (complex(math.cos(O), math.sin(O))) * complex(tx, ty)
                    ekman_flow_x = vxy.real
                    ekman_flow_y = vxy.imag

                # 合并地转流和Ekman流
                total_flow_x = geos_flow_x + ekman_flow_x
                total_flow_y = geos_flow_y + ekman_flow_y

                # 计算总流速和流向
                if total_flow_x != NAN and total_flow_y != NAN:
                    total_flow_v = math.sqrt(total_flow_x ** 2 + total_flow_y ** 2)

                    # 计算表层流场流向angle
                    # 角度定义：以东方向为0度，逆时针方向为正（0-360°）
                    if total_flow_x == 0:
                        if total_flow_y > 0:
                            total_flow_angle = 90
                        elif total_flow_y < 0:
                            total_flow_angle = 270
                        else:
                            total_flow_angle = 0
                    else:
                        total_flow_angle = math.degrees(math.atan(total_flow_y / total_flow_x))
                        if total_flow_y >= 0 and total_flow_x <= 0:
                            total_flow_angle = 180 - abs(total_flow_angle)
                        elif total_flow_y <= 0 and total_flow_x <= 0:
                            total_flow_angle = abs(total_flow_angle) + 180
                        elif total_flow_y <= 0 and total_flow_x >= 0:
                            total_flow_angle = 360 - abs(total_flow_angle)

        # 保存结果
        field_layer.setValue(code, total_flow_x, 0)
        field_layer.setValue(code, ekman_flow_y, 1)
        field_layer.setValue(code, total_flow_v, 2)
        field_layer.setValue(code, total_flow_angle, 3)

        pbar.update(1)

    pbar.close()

    # 保存结果
    output_file = os.path.join(output_file_path,"")
    field_layer.saveToFile(output_file)
    print(f"已保存流场结果: {output_file}")


def findInterRange(code, code_ids, grid_type):
    """
    获取插值所需范围

    :param code:        当前格网ID
    :param code_ids:    所有离散格网的ID
    :param grid_type:   离散格网类型
    """
    result = []
    queue = deque()

    kernel = CreateKernelObj(grid_type)
    queue.append(code)

    # 插值需要至少4个相邻有效格元
    current = code

    # 找第一圈
    fir_query_result = QueryResult()
    kernel.query(DGGSTopoType_RelativeCells, current, fir_query_result)
    query_result_size = fir_query_result.getSize()
    for i in range(query_result_size):
        temp_code = fir_query_result.getGridCodeByIndex(i)
        query_result_one_id = code2Id(temp_code, grid_type)
        if query_result_one_id in code_ids and query_result_one_id not in result:
            result.append(query_result_one_id)

    # 找第二圈
    if len(result) < 4:
        for i in range(query_result_size):
            current = fir_query_result.getGridCodeByIndex(i)
            sec_query_result = QueryResult()
            kernel.query(DGGSTopoType_RelativeCells, current, sec_query_result)
            sec_query_result_size = sec_query_result.getSize()
            for i in range(sec_query_result_size):
                temp_code = sec_query_result.getGridCodeByIndex(i)
                query_result_one_id = code2Id(temp_code, grid_type)
                if query_result_one_id in code_ids and query_result_one_id not in result:
                    result.append(query_result_one_id)

    DestroyKernelObj(kernel)

    return result


def latlon2Coordinate(lat1, lon1, lat2, lon2):
    """
    计算以第一个点(lat1, lon1)为原点，另一个点(lat2, lon2)的局部平面坐标（m）

    :param lat1:    原点纬度
    :param lon1:    原点经度
    :param lat2:    目标纬度
    :param lon2:    目标经度
    :return:        相对坐标
    """
    # 地球表面上同一经线圈上相差1°两点间的距离约为πR / 180 = 111194.926644558737 m
    coordinate_y = (lat2 - lat1) * 111194.926644558737
    # 地球表面上同一纬线圈上相差1°两点间的距离约为πR×cos(纬度) / 180 = 111194.926644 m×cos(纬度)
    coordinate_x = (lon2 - lon1) * 111194.92664458737 * math.cos(math.radians((lat1 + lat2) / 2))

    return coordinate_x, coordinate_y


def getDifferential2(X, Y, Z, x, y, diff_type):
    """
    根据n个三维坐标点拟合二元二次函数方程表达式，返回某切点处偏导数(一阶与二阶偏导数)

    :param X:           数据点横坐标集合
    :param Y:           数据点纵坐标集合
    :param Z:           数据点数据值集合
    :param x:           返回某切点处偏导数，x为该点横坐标
    :param y:           返回某切点处偏导数，y为该点横坐标
    :param diff_type:   返回偏导数类型,包含x', y', x'y', y''
    :return:            拟合后函数(x,y)处偏导数(一阶与二阶偏导数)值
    """
    n = len(X)

    # 求方程系数
    sigma_x = 0
    for i in X: sigma_x += i
    sigma_y = 0
    for i in Y: sigma_y += i
    sigma_z = 0
    for i in Z: sigma_z += i
    sigma_x2 = 0
    for i in X: sigma_x2 += i * i
    sigma_y2 = 0
    for i in Y: sigma_y2 += i * i
    sigma_x3 = 0
    for i in X: sigma_x3 += i * i * i
    sigma_y3 = 0
    for i in Y: sigma_y3 += i * i * i
    sigma_x4 = 0
    for i in X: sigma_x4 += i * i * i * i
    sigma_y4 = 0
    for i in Y: sigma_y4 += i * i * i * i
    sigma_x_y = 0
    for i in range(n):
        sigma_x_y += X[i] * Y[i]

    sigma_x_y2 = 0
    for i in range(n): sigma_x_y2 += X[i] * Y[i] * Y[i]
    sigma_x_y3 = 0
    for i in range(n): sigma_x_y3 += X[i] * Y[i] * Y[i] * Y[i]
    sigma_x2_y = 0
    for i in range(n): sigma_x2_y += X[i] * X[i] * Y[i]
    sigma_x2_y2 = 0
    for i in range(n): sigma_x2_y2 += X[i] * X[i] * Y[i] * Y[i]
    sigma_x3_y = 0
    for i in range(n): sigma_x3_y += X[i] * X[i] * X[i] * Y[i]
    sigma_z_x2 = 0
    for i in range(n): sigma_z_x2 += Z[i] * X[i] * X[i]
    sigma_z_y2 = 0
    for i in range(n): sigma_z_y2 += Z[i] * Y[i] * Y[i]
    sigma_z_x_y = 0
    for i in range(n): sigma_z_x_y += Z[i] * X[i] * Y[i]
    sigma_z_x = 0
    for i in range(n): sigma_z_x += Z[i] * X[i]
    sigma_z_y = 0
    for i in range(n): sigma_z_y += Z[i] * Y[i]

    # 给出对应方程的矩阵形式
    a = np.array([[sigma_x4, sigma_x3_y, sigma_x2_y2, sigma_x3, sigma_x2_y, sigma_x2],
                  [sigma_x3_y, sigma_x2_y2, sigma_x_y3, sigma_x2_y, sigma_x_y2, sigma_x_y],
                  [sigma_x2_y2, sigma_x_y3, sigma_y4, sigma_x_y2, sigma_y3, sigma_y2],
                  [sigma_x3, sigma_x2_y, sigma_x_y2, sigma_x2, sigma_x_y, sigma_x],
                  [sigma_x2_y, sigma_x_y2, sigma_y3, sigma_x_y, sigma_y2, sigma_y],
                  [sigma_x2, sigma_x_y, sigma_y2, sigma_x, sigma_y, n]])
    b = np.array([sigma_z_x2, sigma_z_x_y, sigma_z_y2, sigma_z_x, sigma_z_y, sigma_z])

    # 高斯消元解线性方程
    res = np.linalg.solve(a, b)

    partial_derivatives = {
        "x": 2 * res[0] * x + res[1] * y + res[3],
        "y": res[1] * x + 2 * res[2] * y + res[4],
        "yy": 2 * res[2],
        "xy": res[1]
    }
    partial_derivatives['all'] = [
        partial_derivatives['x'], partial_derivatives['y'],
        partial_derivatives['yy'], partial_derivatives['xy']
    ]

    return partial_derivatives.get(diff_type, 0)


def calCoriolis(lat):
    """
    计算指定纬度lat处的科氏力
    f平面 计算科氏参数 f=2Ωsinφ，Ω是地球自转角速度(单位：角度/秒)，φ是纬度

    :param lat: 纬度
    :return:    科氏力
    """
    # 科式参数
    f = 2 * (2 * math.pi / 24 / 3600) * math.sin(math.radians(lat))  
    return f


def calSwayingCoefficient(wind_speed):
    """
    计算Ekman流中所需的摇曳系数

    :param wind_speed:  风速值
    :return:            指定风速下的摇曳系数
    """

    if wind_speed == 0:         #  处理风速为极小值的特殊情况
        wind_speed = 1e-8
        c = 0.62 + 1.56 / wind_speed 
    elif wind_speed < 3:        #  根据当风速小于3时，计算摇曳系数公式的常数项为：a = 1.56; b = 0.62
        c = 0.62 + 1.56 / wind_speed
    elif 3 <= wind_speed <= 10: #  当风速在3到10之间时，摇曳系数固定为1.14
        c = 1.14
    else:                       #  当风速大于10时，计算摇曳系数的公式的常数项为：a = 0.49; b = 0.065
        c = 0.49 + 0.065 * wind_speed
    c = c / 1000.0              

    return c


def calEquatorEkman(lat, tx, ty):
    """
    计算Ekman流的通用公式

    :param lat:    纬度
    :param tx:     x方向风应力
    :param ty:     y方向凤应力
    :return (x方向ekman流分量, y方向ekman流分量 )
    """
    f = calCoriolis(lat)
    seawater_density = 1.02 * (10 ** 3) # 海水密度kg/m^3
    friction = 2.15 * (10 ** -4)        # 摩擦系数 m/s
    mix_layer_depth = 32.5              # 混合层深度 m

    B = 1.0 / (seawater_density * (friction ** 2 + (f * mix_layer_depth) ** 2) ** 0.5)  # B为中间参数
    O = math.atan(f * mix_layer_depth / friction)  # θ是流向相对于风向的旋转角(角度值)   #!!!
    # O = abs(O)

    # Ekman流速度 vx vy
    vxy = B * (complex(math.cos(O), math.sin(O))) * (tx + ty)
    vx = vxy.real
    vy = vxy.imag

    return (vx, vy)


def readOceanGridCodeid(ocean_grid_file):
    """
    读取海洋格网对应的文件中codeId属性值

    :param ocean_grid_file:     海洋格网文件路径
    :return code_ids:           海洋格网的codeId列表
    """
    # 创建空列表，用于储存codeId
    code_ids = []

    # 导入海洋格网
    ocean_grid = DGGSFieldLayerInt()
    ocean_grid.loadFromFile(ocean_grid_file)

    ocean_band = 0
    for base_partition in range(1,11):
        for row in range(rowMaxSize(level)):
            for col in range(fColMaxSize(level, grid_type)):
                code = DGGSElementCode()
                code.elementType = element_type
                code.level = level
                code.basePartition = base_partition
                code.row = row
                code.col = col
                value = ocean_grid.getValue(code, ocean_band)

                # 判断是否为海洋，是海洋的格网则添加id到code_ids中
                if value == 1:
                    # 获取格元的编码
                    code_id = code2Id(code, grid_type)
                    code_ids.append(code_id)

    return code_ids


def saveUvFile(year, month, day, 
                level, grid_type, element_type,
                temp_path, adt_input_path, wind_input_path, output_path, ocean_grid_file):
    """
    保存uv分量的gridtiff格式数据

    :param year:                年
    :param month:               月
    :param day:                 日
    :param level:               格网等级
    :param grid_type:           格网类型
    :param element_type:        元素类型
    :param temp_path:           临时路径
    :param adt_input_path:      ADT输入路径
    :param wind_input_path:     风场输入路径
    :param output_path:         输出路径
    :param ocean_grid_file:     海洋格网文件
    """
    # 生成索引
    adt_vars = ['adt']
    wind_vars = ['eastward_wind', 'northward_wind']
    ocean_grid_ids = readOceanGridCodeid(ocean_grid_file)
    
    # 根据day参数确定日期格式
    if day != '':
        date_str = f"{year}{month}{day}"
        output_file = os.path.join(output_path, f"{year}{month}{day}")
        print(f"处理单日数据: {year}-{month}-{day}")
    else:
        date_str = f"{year}{month}"
        output_file = os.path.join(output_path, f"{year}{month}")
        print(f"处理月度数据: {year}-{month}")
    
    # 处理ADT数据
    print(f"处理ADT数据: {adt_input_path}")
    createIndexPath(
        adt_vars, year, month, day,
        adt_input_path, temp_path,
        level, grid_type, element_type, ocean_grid_ids
    )
    
    # 处理Wind数据
    print(f"处理Wind数据: {wind_input_path}")
    createIndexPath(
        wind_vars, year, month, day,
        wind_input_path, temp_path,
        level, grid_type, element_type, ocean_grid_ids
    )
    
    # 准备输入文件
    input_files = [
        os.path.join(temp_path, f"{date_str}_eastward_wind.csv"),
        os.path.join(temp_path, f"{date_str}_northward_wind.csv"),
        os.path.join(temp_path, f"{date_str}_adt.csv")
    ]
    
    # 计算流场
    print("计算海洋流场:")
    calOceanFlow(year, month, day, input_files, output_file, level, grid_type)
    
    print(f"处理完成: {output_file}")

def main():
    
    months = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]   # 一年的月份

    if time_resolution == 'daily':    # 时间分辨率是日
        years = [2015, 2016]  # 年份
        for year in years:
            for month in months:
                _, days_in_month = calendar.monthrange(year, month)
                for day in range(1, days_in_month + 1):
                    year_str = str(year)
                    month_str = f'{month:02d}'
                    day_str = f'{day:02d}'

                    os.makedirs(os.path.join(daily_uv_output_path, f"{year_str}{month_str}{day_str}"))
                    saveUvFile(year_str, month_str, day_str,
                                level, grid_type, element_type,
                                daily_temp_output_path,daily_input_adt_path, daily_input_wind_path,daily_uv_output_path, ocean_grid_path)

    
    elif time_resolution == 'monthly':    # 时间分辨率是月
        years = [2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020]     # 年份
        day = ''
        for year in years:
            for month in months:
                year_str = str(year)
                month_str = f'{month:02d}'
                
                os.makedirs(os.path.join(monthly_uv_output_path, f"{year_str}{month_str}"))
                saveUvFile(year_str, month_str, day,
                            level, grid_type, element_type, 
                            monthly_temp_output_path,monthly_input_adt_path, monthly_input_wind_path, monthly_uv_output_path, ocean_grid_path)

    else:
        print("未知的时间分辨率")
                

if __name__ == "__main__":
    start_time = time.time()
    main()
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"程序执行时间: {execution_time} 秒")